from mmengine.config import read_base
with read_base():
    from .mgam import *

from mgamdata.models.SegFormer3D import (
    SegFormer3D_Encoder_MM, 
    SegFormer3D_Decoder_MM)
from mgamdata.models.ReconSelfSup import Reconstructor, ReconHead
from mgamdata.mm.mmseg_Dev3D import DiceLoss_3D


embed_dims = 32
embed_dims_SegFormer3D = [
    embed_dims, 2*embed_dims, 5*embed_dims, 8*embed_dims]
head_embed_dims = embed_dims_SegFormer3D[-1]//2

model = dict(
    type = Reconstructor,
    recon_channels=1,
    encoder=dict(
        type=SegFormer3D_Encoder_MM, 
        in_channels=in_channels, # type: ignore
        embed_dims=embed_dims_SegFormer3D, 
    ), 
    neck=None,
    decoder=dict(
        type=SegFormer3D_Decoder_MM, 
        embed_dims=embed_dims_SegFormer3D, 
        head_embed_dims=head_embed_dims, 
        num_classes=head_embed_dims,
        loss_decode=dict(
            type=DiceLoss_3D, 
            use_sigmoid=False, 
            ignore_1st_index=False, 
            batch_z=4, 
            # NOTE Severe performance overhead when not being set to None.
            # NOTE Prefer using `ignore_1st_index`.
            # NOTE Invalid Class (Defaults to the last class) has been masked out during preprocess.
            ignore_index=None, 
        ), 
    ),
    head=dict(
        type=ReconHead, 
        model_out_channels=head_embed_dims, 
        recon_channels=1,
        dim='3d',
    ),
    test_cfg=dict(
        mode="slide",
        crop_size=size,
        slide_accumulate_device='cuda',
        stride=[i//2 for i in size]
    )
)
